import asyncio
from datetime import datetime
from typing import Dict, Optional, List, Any
import logging
import json
import re

# LangChain核心查询依赖（适配v0.2+）
from langchain_milvus import Milvus
from langchain_ollama import OllamaEmbeddings,OllamaLLM
from langchain.chains import RetrievalQA
from langchain.prompts import PromptTemplate
from langchain_core.documents import Document
from pymilvus import Collection, connections, utility
from app.config import settings
from langchain_community.retrievers import BM25Retriever
from langchain_core.runnables import RunnablePassthrough
from langchain_core.output_parsers import StrOutputParser
from app.excepetions.rag_exception import RAGException

# 自定义异常（仅用于查询流程错误）
class RAGQueryException(Exception):
    def __init__(self, code: int, message: str):
        self.code = code
        self.message = message
        super().__init__(f"[{code}] {message}")


# 初始化日志（仅保留查询相关日志）
logging.basicConfig(
    level=logging.INFO,
    format="%(asctime)s - NRS2002-RAG-Query - %(levelname)s - %(message)s"
)
logger = logging.getLogger("RAGQueryService")


class RAGQueryService:
    def __init__(
            self,
            milvus_host: str = settings.MILVUS_HOST,  # 改用host/port方式连接
            milvus_port: str = settings.MILVUS_PORT,
            milvus_token: Optional[str] = None,  # Milvus认证令牌（如无则为None）
            collection_name: str = settings.MILVUS_DEFAULT_COLLECTION,  # Milvus集合名称
            embedding_model: str = settings.EMBEDDING_MODEL,
            llm_model: str = settings.LLM_MODEL,
            ollama_base_url: str = settings.OLLAMA_BASE_URL,
            dim: int = 1024,  # 嵌入向量维度，根据模型调整（bge-m3为1024）
            require_data: bool = True,  # 新增：是否要求集合必须有数据
            database_name: str = "test_db"
    ):
        """初始化查询必需组件：嵌入模型、Milvus向量库、LLM、QA链"""
        # 打印Ollama基础URL用于调试
        print(f"self.ollama_base_url:{ollama_base_url}")
        logger.info(f"Ollama基础URL: {ollama_base_url}")

        self.milvus_host = milvus_host
        self.milvus_port = milvus_port
        self.milvus_token = milvus_token
        self.collection_name = collection_name
        self.embedding_model = embedding_model
        self.llm_model = llm_model
        self.ollama_base_url = ollama_base_url
        self.dim = dim  # 向量维度
        self.require_data = require_data  # 控制是否检查数据存在性11
        self.database_name = database_name
        # 初始化核心组件（顺序不可变）
        self.embeddings = self._init_embeddings()
        self.vector_store = self._init_vector_store()
        self.llm = self._init_llm()
        self.qa_chain = self._init_qa_chain()
        self.bm25_retriever = self._init_bm25_retriever()

    def _init_embeddings(self) -> OllamaEmbeddings:
        """初始化嵌入模型（验证可用性）"""
        try:
            embeddings = OllamaEmbeddings(
                model=self.embedding_model,
                base_url=self.ollama_base_url,
                client_kwargs={"timeout": 600.0}  # 增加超时设置
            )
            logger.info(f"embeddings loaded: {embeddings}:{self.embedding_model},{self.ollama_base_url}")

            # 验证嵌入输出有效性并获取实际维度
            test_emb = embeddings.embed_query("NRS2002营养风险筛查")
            if not (isinstance(test_emb, list) and len(test_emb) > 0 and isinstance(test_emb[0], float)):
                raise RAGQueryException(503, "嵌入模型返回无效向量")

            # 检查维度是否匹配配置
            actual_dim = len(test_emb)
            if actual_dim != self.dim:
                logger.warning(f"嵌入维度不匹配，配置维度: {self.dim}, 实际维度: {actual_dim}，将使用实际维度")
                self.dim = actual_dim

            logger.info(f"嵌入模型就绪：{self.embedding_model}（向量维度：{self.dim}）")
            return embeddings
        except Exception as e:
            raise RAGQueryException(500, f"嵌入模型初始化失败：{str(e)}")

    def _init_vector_store(self) -> Milvus:
        """初始化Milvus向量库（修复异步连接问题）"""
        try:
            # 确保在当前线程有事件循环
            try:
                asyncio.get_event_loop()
            except RuntimeError:
                # 没有事件循环则创建新的
                loop = asyncio.new_event_loop()
                asyncio.set_event_loop(loop)
            # 断开现有连接（如果存在）
            if connections.has_connection("default"):
                connections.disconnect("default")

            # 使用同步方式连接（避免异步事件循环问题）
            connections.connect(
                host=self.milvus_host,
                port=self.milvus_port,
                token=self.milvus_token,
                alias="default",
                pool_type="CPU"  ,# 强制使用同步连接池（非异步）


            )

            # 检查集合是否存在
            if not utility.has_collection(self.collection_name):
                raise RAGQueryException(404, f"Milvus集合'{self.collection_name}'不存在（需先创建并入库NRS2002规则文档）")

            # 连接Milvus集合（使用host/port而非uri）
            vector_store = Milvus(
                embedding_function=self.embeddings,
                connection_args={
                    "host": self.milvus_host,
                    "port": self.milvus_port,
                    "token": self.milvus_token,
                    # 额外保险：部分版本需通过connection_args禁用异步
                    "enable_async": False
                },
                collection_name=self.collection_name,
                drop_old=False,  # 不删除旧集合
            )

            # 检查集合中的实体数量（仅当require_data为True时）
            if self.require_data:
                coll = Collection(self.collection_name)
                coll.load()  # 确保集合已加载
                count = coll.num_entities
                if count == 0:
                    raise RAGQueryException(400, f"Milvus集合'{self.collection_name}'无数据（需先入库NRS2002规则文档）")
                logger.info(f"Milvus集合'{self.collection_name}'包含{count}条数据")

            logger.info(f"Milvus向量库就绪：{self.collection_name}")
            return vector_store
        except RAGQueryException:
            # 重新抛出已知异常
            raise
        except Exception as e:
            raise RAGQueryException(500, f"Milvus初始化失败：{str(e)}")

    def _init_llm(self) -> OllamaLLM:
        """初始化LLM（修复类名错误并验证可用性）"""
        try:
            # 注意：正确的类名是Ollama而非OllamaLLM
            llm = OllamaLLM(
                model=self.llm_model,
                base_url=self.ollama_base_url,
                temperature=0.1,  # 低温度确保规则严格执行
                client_kwargs={"timeout": 60.0}  # 增加超时设置
            )

            # 测试LLM响应
            test_resp = llm.invoke("仅返回'pong'")
            if "pong" not in test_resp.strip().lower():
                raise RAGQueryException(503, f"LLM测试失败：返回'{test_resp[:30]}'")

            logger.info(f"LLM就绪：{self.llm_model}")
            return llm
        except Exception as e:
            error_msg = str(e)
            if "connection refused" in error_msg.lower():
                error_msg = f"Ollama服务连接失败（地址：{self.ollama_base_url}）"
            elif "model not found" in error_msg.lower():
                error_msg = f"LLM模型不存在（需执行'ollama pull {self.llm_model}'）"
            raise RAGQueryException(503, error_msg)

    def _init_qa_chain(self) -> RetrievalQA:
        """初始化QA链（绑定NRS2002规则Prompt，确保输出带依据）"""
        nrs2002_prompt = PromptTemplate(
            template="""
任务：根据NRS2002营养风险筛查规则，基于参考上下文计算患者评分，输出JSON（含评分和依据说明）。

NRS2002核心规则（必须严格遵守）：
1. 营养状态受损（0-3分）：
   - 0分：BMI≥20.5且体重无下降且无进食困难
   - 1分：BMI18.5-20.4/近3月体重降3%-5%/进食减25%-50%
   - 2分：BMI＜18.5/近3月体重降5%-10%/进食减50%-75%
   - 3分：近3月体重降＞10%/进食减＞75%/BMI＜18.5+重病
2. 疾病严重程度（0-3分）：
   - 0分：良性疾病+正常进食（如稳定期慢性病）
   - 1分：慢病急性发作+需卧床（如COPD急性加重）
   - 2分：大手术/中风/ICU+需人工营养
   - 3分：大面积烧伤/多器官衰竭+紧急营养
3. 年龄（0-1分）：≥70岁得1分，否则0分
4. 总分=三部分之和（缺失参数默认0分，总分范围0-7分）

要求：
1. 先分析患者信息匹配哪条规则，再计算各维度分数和总分；
2. 输出JSON必须包含"score"（总分）、"nutritional_impairment"（营养受损分）、"disease_severity"（疾病严重度分）、"age"（年龄分）、"basis"（评分依据，说明匹配的规则条款）；
3. 仅输出JSON，无多余文字。

用户问题（患者信息）：{question}
参考上下文（NRS2002规则片段）：{context}

输出格式示例：
{{
  "score": 2,
  "nutritional_impairment": 1,
  "disease_severity": 1,
  "age": 0,
  "basis": "1.营养受损：BMI19.2（18.5-20.4）→1分；2.疾病严重度：COPD急性加重→1分；3.年龄65岁＜70→0分；总分1+1+0=2分"
}}
""",
            input_variables=["context", "question"]
        )

        return RetrievalQA.from_chain_type(
            llm=self.llm,
            chain_type="stuff",
            retriever=self.vector_store.as_retriever(search_kwargs={"k": 2}),  # 取最相关的2个规则片段
            chain_type_kwargs={"prompt": nrs2002_prompt},
            return_source_documents=True  # 返回源文档作为验证依据
        )

    def query_score(self, user_question: str, file_id: Optional[str] = None,collection_name: Optional[str] = None) -> Dict[str, Any]:
        """核心查询方法：输入患者信息，返回NRS2002评分、依据说明和源文档片段"""
        # 1. 输入校验
        if not user_question.strip():
            raise RAGQueryException(400, "患者信息不能为空（需包含BMI、体重变化、疾病状态、年龄等关键信息）")

        try:
            # 2. 构建检索过滤条件
            search_kwargs = {"k": 2}
            if file_id:
                # 处理file_id中的单引号
                safe_file_id = str(file_id).replace("'", "''")
                search_kwargs["expr"] = f"file_id == '{safe_file_id}'"
                logger.info(f"Milvus过滤条件expr: {search_kwargs['expr']}")

                # 修复计数逻辑：使用Collection的query方法
                selected_collection = collection_name if collection_name else self.collection_name
                coll = Collection(selected_collection)
                if not coll.is_empty:
                    coll.load()

                # 直接查询判断是否存在匹配数据
                res = coll.query(
                    expr=search_kwargs["expr"],
                    output_fields=["pk"],
                    limit=1
                )
                count = len(res)
                logger.info(f"符合条件的实体数量: {count}")

                if count == 0:
                    raise RAGQueryException(404, f"指定file_id={file_id}无匹配的NRS2002规则数据")

            self.qa_chain.retriever.search_kwargs = search_kwargs

            # 3. 执行RAG查询
            logger.info(f"执行查询：{user_question[:50]}...")
            result = self.qa_chain.invoke({"query": user_question.strip()})

            # 4. 解析并验证LLM输出
            json_match = re.search(r'\{[\s\S]*?\}', result["result"].strip())
            if not json_match:
                raise RAGQueryException(500, f"LLM未输出有效JSON：{result['result'][:100]}")

            try:
                score_data = json.loads(json_match.group())
            except json.JSONDecodeError as e:
                raise RAGQueryException(500, f"JSON解析失败：{str(e)}（原始输出：{json_match.group()[:50]}）")

            # 验证必需字段
            required_fields = ["score", "nutritional_impairment", "disease_severity", "age", "basis"]
            for field in required_fields:
                if field not in score_data:
                    raise RAGQueryException(500, f"评分结果缺失必需字段：{field}")
                if field != "basis" and not isinstance(score_data[field], int):
                    raise RAGQueryException(500, f"字段{field}必须为整数（当前：{score_data[field]}）")

            # 验证分数范围
            if not (0 <= score_data["score"] <= 7):
                raise RAGQueryException(500, f"总分必须在0-7分之间（当前：{score_data['score']}）")
            for field in ["nutritional_impairment", "disease_severity"]:
                if not (0 <= score_data[field] <= 3):
                    raise RAGQueryException(500, f"{field}必须在0-3分之间（当前：{score_data[field]}）")
            if score_data["age"] not in [0, 1]:
                raise RAGQueryException(500, f"年龄分必须为0或1（当前：{score_data['age']}）")

            # 5. 整理源文档依据
            source_basis = []
            for doc in result["source_documents"]:
                source_basis.append({
                    "id": doc.metadata.get("id", "unknown"),
                    "file_name": doc.metadata.get("file_name", "unknown"),
                    "file_id": doc.metadata.get("file_id", "unknown"),
                    "rule_fragment": doc.page_content.strip()[:200] + "..." if len(
                        doc.page_content) > 200 else doc.page_content.strip()
                })

            # 6. 返回最终结果
            return {
                "code": 200,
                "message": "查询成功",
                "score_result": score_data,
                "source_basis": source_basis
            }

        except RAGQueryException as e:
            logger.error(f"查询失败：{e.message}")
            raise
        except Exception as e:
            raise RAGQueryException(500, f"查询流程异常：{str(e)}")


    def _init_bm25_retriever(self) -> BM25Retriever:
        """初始化BM25检索器（从Milvus加载已存储的文档）"""
        try:
            # 1. 从Milvus读取所有文档（需确保Milvus已存储文档，可在process后调用）
            all_docs = self.vector_store.similarity_search("")  # 空查询获取所有文档（Milvus特性）
            if not all_docs:
                logger.warning("BM25初始化：Milvus中暂无文档，后续需处理文档后重新加载")
                return BM25Retriever.from_documents([])

            # 2. 用Jieba分词处理文档（适配中文，提升BM25效果）
            def jieba_tokenize(text: str) -> List[str]:
                import jieba
                return list(jieba.cut(text))

            # 3. 初始化BM25检索器（指定中文分词器）
            bm25_retriever = BM25Retriever.from_documents(
                all_docs,
                tokenizer=jieba_tokenize, # 关键：用Jieba替换默认英文分词


            )
            # 设置BM25返回结果数（可调整，与向量检索一致）
            bm25_retriever.k = 5
            logger.info("BM25检索器初始化完成（基于Jieba中文分词）")
            return bm25_retriever
        except Exception as e:
            raise RAGException(500, f"BM25检索器初始化失败：{str(e)}")

    def _vector_search(self, query: str, k: int = 5) -> List[Document]:
        """纯向量检索（Milvus）：适合语义关联查询"""
        try:
            # 调用Milvus向量检索，返回Top-k结果
            results = self.vector_store.similarity_search(
                query=query,
                k=k,
                # 可选：过滤条件（如只检索营养测评相关文档，基于metadata的file_type字段）
                filter={"file_type": "nutrition_assessment"}
            )
            logger.info(f"向量检索完成：查询='{query}'，返回{len(results)}条结果")
            return results
        except Exception as e:
            raise RAGException(500, f"向量检索失败：{str(e)}")

    def _bm25_search(self, query: str, k: int = 5) -> List[Document]:
        """纯BM25检索：适合关键词精确查询（如营养测评指标）"""
        try:
            # 更新BM25检索器的返回数
            self.bm25_retriever.k = k
            # 执行BM25检索
            results = self.bm25_retriever.get_relevant_documents(query=query)
            logger.info(f"BM25检索完成：查询='{query}'，返回{len(results)}条结果")
            return results
        except Exception as e:
            raise RAGException(500, f"BM25检索失败：{str(e)}")

    def hybrid_search(self, query: str, k: int = 5, vector_weight: float = 0.6, bm25_weight: float = 0.4) -> List[
        Document]:
        """
        混合检索（向量检索+BM25）：适配营养测评场景的精准检索
        Args:
            query: 检索查询（如“糖尿病患者NRS2002评分标准”）
            k: 最终返回结果数
            vector_weight: 向量检索结果权重（建议0.5-0.7）
            bm25_weight: BM25检索结果权重（建议0.3-0.5）
        Returns:
            加权排序后的文档列表
        """
        try:
            # 1. 分别执行两种检索（获取带得分的结果）
            # 向量检索：用similarity_search_with_score获取余弦相似度（0-1，越大越相关）
            vector_docs_with_score = self.vector_store.similarity_search_with_score(query, k=k * 2)  # 多取一倍用于去重
            # BM25检索：用get_relevant_documents_with_score获取BM25得分（无固定范围，越大越相关）
            bm25_docs_with_score = self.bm25_retriever.get_relevant_documents_with_score(query, k=k * 2)

            # 2. 标准化BM25得分（转为0-1范围，与向量得分对齐）
            bm25_scores = [score for _, score in bm25_docs_with_score]
            max_bm25 = max(bm25_scores) if bm25_scores else 1
            min_bm25 = min(bm25_scores) if bm25_scores else 0
            normalized_bm25 = [(doc, (score - min_bm25) / (max_bm25 - min_bm25)) for doc, score in bm25_docs_with_score]

            # 3. 合并结果（去重，按doc.page_content唯一标识）
            doc_map = {}
            # 加入向量检索结果
            for doc, vec_score in vector_docs_with_score:
                key = doc.page_content  # 用文档内容去重（也可改用file_id+chunk_id）
                if key not in doc_map:
                    doc_map[key] = {"doc": doc, "vec_score": vec_score, "bm25_score": 0.0}
                else:
                    doc_map[key]["vec_score"] = max(doc_map[key]["vec_score"], vec_score)  # 取最高向量得分
            # 加入BM25检索结果
            for doc, bm25_score in normalized_bm25:
                key = doc.page_content
                if key not in doc_map:
                    doc_map[key] = {"doc": doc, "vec_score": 0.0, "bm25_score": bm25_score}
                else:
                    doc_map[key]["bm25_score"] = max(doc_map[key]["bm25_score"], bm25_score)  # 取最高BM25得分

            # 4. 加权计算最终得分，排序后取Top-k
            sorted_docs = sorted(
                doc_map.values(),
                key=lambda x: (x["vec_score"] * vector_weight) + (x["bm25_score"] * bm25_weight),
                reverse=True  # 降序：得分越高越相关
            )[:k]

            # 5. 提取最终文档列表
            final_results = [item["doc"] for item in sorted_docs]
            logger.info(f"混合检索完成：查询='{query}'，返回{len(final_results)}条结果（去重后）")
            return final_results
        except Exception as e:
            raise RAGException(500, f"混合检索失败：{str(e)}")

    def nutrition_assessment_qa(self, query: str, k: int = 5) -> Dict[str, Any]:
        """
        营养测评场景问答：混合检索+LLM生成结构化回答
        Args:
            query: 用户查询（如“NRS2002评分对糖尿病患者的标准是什么？”）
        Returns:
            包含检索结果和生成回答的字典
        """
        try:
            # 1. 执行混合检索（核心：获取相关文档）
            retrieved_docs = self.hybrid_search(query, k=k)
            if not retrieved_docs:
                return {
                    "query": query,
                    "has_result": False,
                    "answer": "未找到相关营养测评资料，请补充查询关键词（如NRS2002、MUST等）",
                    "retrieved_docs": []
                }

            # 2. 构建LLM提示词（注入检索结果，要求结构化输出）
            prompt_template = """
            你是营养测评专家，基于以下参考资料，精准回答用户问题：
            {context}

            用户问题：{query}

            回答要求：
            1. 只使用参考资料中的信息，不编造内容；
            2. 结构清晰（分点说明，如适用人群、评分步骤、判断标准）；
            3. 针对营养测评场景，突出关键指标（如BMI、疾病严重程度、进食情况）；
            4. 若资料中无明确答案，需说明“参考资料中未提及该细节”。
            """

            # 3. 组装上下文（检索结果的文档内容）
            context = "\n\n".join([f"【资料{idx + 1}】{doc.page_content}" for idx, doc in enumerate(retrieved_docs)])

            # 4. 调用LLM生成回答
            chain = (
                    {"context": RunnablePassthrough(), "query": RunnablePassthrough()}
                    | prompt_template
                    | self.llm
                    | StrOutputParser()
            )
            answer = chain.invoke({"context": context, "query": query})

            # 5. 整理检索结果元数据（返回给用户查看来源）
            retrieved_metadata = [
                {
                    "file_id": doc.metadata.get("file_id", "unknown"),
                    "file_name": doc.metadata.get("file_name", "unknown"),
                    "chunk_id": doc.metadata.get("chunk_id", "unknown"),
                    "source": doc.metadata.get("source", "unknown")
                }
                for doc in retrieved_docs
            ]

            return {
                "query": query,
                "has_result": True,
                "answer": answer,
                "retrieved_docs_count": len(retrieved_docs),
                "retrieved_docs_metadata": retrieved_metadata
            }
        except Exception as e:
            raise RAGException(500, f"营养测评问答失败：{str(e)}")

# # 测试入口
# if __name__ == "__main__":
#     try:
#         logger.info("=" * 50)
#         logger.info("初始化NRS2002 RAG查询服务（基于Milvus）")
#         rag_query_service = RAGQueryService(
#             milvus_host="localhost",  # 使用host/port连接方式
#             milvus_port="19530",
#             milvus_token=None,
#             collection_name="nrs2002_collection",
#             embedding_model="bge-m3:latest",
#             llm_model="qwen3:8b",
#             ollama_base_url="http://127.0.0.1:31434",  # 匹配实际Ollama地址
#             dim=1024,
#             require_data=True  # 首次运行确保有数据，可设为False进行调试
#         )
#         logger.info("查询服务初始化成功")
#
#         # 测试用例
#         test_cases = [
#             {
#                 "id": 1,
#                 "question": "患者女性，65岁，BMI 19.2，2型糖尿病稳定期（正常进食），近3个月体重从65kg降至62kg（下降4.6%），无进食困难"
#             }
#         ]
#
#         # 执行测试
#         for case in test_cases:
#             logger.info(f"\n{'=' * 60}")
#             logger.info(f"测试用例{case['id']}：{case['question']}")
#             logger.info(f"{'=' * 60}")
#
#             result = rag_query_service.query_score(case["question"])
#             print(f"\n【测试用例{case['id']} - 查询结果】")
#             print(f"1. 评分结果：\n{json.dumps(result['score_result'], ensure_ascii=False, indent=2)}")
#
#     except RAGQueryException as e:
#         logger.error(f"\n测试失败：[{e.code}] {e.message}", exc_info=True)
#         exit(1)
#     except Exception as e:
#         logger.error(f"\n未知错误：{str(e)}", exc_info=True)
#         exit(1)
